package net.lab1024.sa.ai.stream.module.business.ai.manager;

import com.alibaba.fastjson.JSON;
import lombok.extern.slf4j.Slf4j;

import net.lab1024.sa.ai.stream.module.business.ai.entity.AiModelConfigEntity;
import net.lab1024.sa.ai.stream.module.business.ai.integration.AiModelAdapter;
import net.lab1024.sa.ai.stream.module.business.ai.service.AiModelConfigService;
import net.lab1024.sa.ai.stream.module.business.ai.service.AiRolePromptService;
import net.lab1024.sa.ai.stream.module.business.ai.entity.AiRolePromptEntity;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;
import java.util.HashMap;
import java.util.Map;
import java.util.List;
import reactor.core.publisher.Flux;

@Slf4j
@Component
public class AiModelManager {

    @Autowired
    private AiModelConfigService aiModelConfigService;
    @Autowired
    private AiRolePromptService aiRolePromptService;

    private final Map<String, AiModelAdapter> adapterMap = new HashMap<>();

    @Autowired
    public AiModelManager(List<AiModelAdapter> adapters) {
        for (AiModelAdapter adapter : adapters) {
            adapterMap.put(adapter.getModelName().toLowerCase(), adapter);
        }
    }

    public String chat(String modelName, String prompt, Map<String, Object> params, Long roleId, List<String> promptIds) {
        AiModelConfigEntity config = aiModelConfigService.getByModelName(modelName);
        if (config == null) {
            throw new RuntimeException("模型配置不存在: " + modelName);
        }
        Map<String, Object> mergedParams = mergeParams(config.getParamsJson(), params);
        // 组装 messages
        List<Map<String, String>> messages = new java.util.ArrayList<>();
        if (promptIds != null && !promptIds.isEmpty()) {
            List<AiRolePromptEntity> prompts = aiRolePromptService.listByAliasList(roleId, promptIds);
            for (AiRolePromptEntity p : prompts) {
                messages.add(java.util.Map.of("role", "system", "content", p.getPrompt()));
            }
        }
        messages.add(java.util.Map.of("role", "user", "content", prompt));
        mergedParams.put("messages", messages);
        AiModelAdapter adapter = getAdapter(modelName);
        return adapter.chat(null, mergedParams, config);
    }

    public Flux<String> chatStream(String modelName, String prompt, Map<String, Object> params, Long roleId, List<String> promptIds) {
        AiModelConfigEntity config = aiModelConfigService.getByModelName(modelName);
        if (config == null) {
            throw new RuntimeException("模型配置不存在: " + modelName);
        }
        Map<String, Object> mergedParams = mergeParams(config.getParamsJson(), params);
        // 组装 messages
        List<Map<String, String>> messages = new java.util.ArrayList<>();
        if (promptIds != null && !promptIds.isEmpty()) {
            List<AiRolePromptEntity> prompts = aiRolePromptService.listByAliasList(roleId, promptIds);
            for (AiRolePromptEntity p : prompts) {
                // log.info("添加提示词: {}", p.getPrompt());  
                messages.add(java.util.Map.of("role", "system", "content", p.getPrompt()));
            }
        }
        messages.add(java.util.Map.of("role", "user", "content", prompt));
        mergedParams.put("messages", messages);
        AiModelAdapter adapter = getAdapter(modelName);
        return adapter.chatStream(null, mergedParams, config);
    }

    private Map<String, Object> mergeParams(String paramsJson, Map<String, Object> userParams) {
        Map<String, Object> defaultParams = new HashMap<>();
        if (paramsJson != null && !paramsJson.isEmpty()) {
            defaultParams = JSON.parseObject(paramsJson, Map.class);
        }
        if (userParams != null) {
            defaultParams.putAll(userParams); // 用户参数优先生效
        }
        return defaultParams;
    }

    private AiModelAdapter getAdapter(String modelName) {
        AiModelAdapter adapter = adapterMap.get(modelName.toLowerCase());
        if (adapter == null) {
            throw new RuntimeException("暂不支持的模型: " + modelName);
        }
        return adapter;
    }
} 